#include <jni.h>
#include <string>
#include <android/bitmap.h>

#include "print_log.h"
#include "ncnn_instance.h"

extern "C" JNIEXPORT jlong

JNICALL
Java_com_lanyeo_liver_util_NCNN_resourceTaskCreate(
        JNIEnv *env,
        jobject, /* this */
        jstring jModelDir) {

    const char* modelDir = env->GetStringUTFChars(jModelDir, JNI_FALSE);
    LOGD("resourceTaskCreate dir:%s", modelDir);
    auto task = new NCNNInstance();
    int ret = task->Init(modelDir);
    if (ret != 0) {
        LOGE("resourceTaskCreate fail. ret:%d", ret);
        delete task;
        return 0;
    }
    return reinterpret_cast<jlong>(task);
}

extern "C" JNIEXPORT jint

JNICALL
Java_com_lanyeo_liver_util_NCNN_resourceTaskDetect(
        JNIEnv *env,
        jobject /* this */,
        jlong jTaskObjAddress,
        jobject jBitmap,
        jobject jObjects /* std::vector<Object> <-> ArrayList<Object>*/
) {
    if (jTaskObjAddress == 0 || jBitmap == 0) {
        LOGE("resourceTaskDetect fail. param error!");
        return -1;
    }

    void* bitmapPixels = nullptr;
    AndroidBitmapInfo bitmapInfo;
    AndroidBitmap_getInfo(env, jBitmap, &bitmapInfo);
    AndroidBitmap_lockPixels(env, jBitmap, &bitmapPixels);
    if (bitmapPixels == nullptr) {
        return -1;
    }
    if (bitmapInfo.width < bitmapInfo.height) {
        AndroidBitmap_unlockPixels(env, jBitmap);
        return -1;
    }
    cv::Mat img = cv::Mat(bitmapInfo.height, bitmapInfo.width, CV_8UC4, bitmapPixels);

    auto* task = (NCNNInstance*)jTaskObjAddress;

    int ret = task->Detect(img);
    AndroidBitmap_unlockPixels(env, jBitmap);
    // LOGD("resourceTaskDetect size:%zu ret:%d", task->objects().size(), ret);
    if (ret != 0) {
        return ret;
    }

    // new一个ArrayList对象
    jclass class_arraylist=env->FindClass("java/util/ArrayList");
    jmethodID arraylist_add_method= env->GetMethodID(class_arraylist,"add","(Ljava/lang/Object;)Z");

    // Rect对象
    jclass class_rect=env->FindClass("org/opencv/core/Rect");
    jmethodID construct_rect=env->GetMethodID(class_rect, "<init>","()V");
    jfieldID rect_w=env->GetFieldID(class_rect, "width","I");
    jfieldID rect_h=env->GetFieldID(class_rect, "height","I");
    jfieldID rect_x=env->GetFieldID(class_rect, "x","I");
    jfieldID rect_y=env->GetFieldID(class_rect, "y","I");

    // Object对象
    jclass class_object=env->FindClass("com/lanyeo/liver/util/Object");
    jmethodID construct_object=env->GetMethodID(class_object, "<init>","()V");
    jfieldID object_label=env->GetFieldID(class_object, "label","I");
    jfieldID object_rect=env->GetFieldID(class_object, "rect","Lorg/opencv/core/Rect;");

    for (const Object& obj: task->objects()) {
        jobject obj_rect = env->NewObject(class_rect, construct_rect);
        env->SetIntField(obj_rect, rect_w, obj.rect.width);
        env->SetIntField(obj_rect, rect_h, obj.rect.height);
        env->SetIntField(obj_rect, rect_x, obj.rect.x);
        env->SetIntField(obj_rect, rect_y, obj.rect.y);

        jobject jObj = env->NewObject(class_object, construct_object);
        env->SetIntField(jObj, object_label, obj.label);
        env->SetObjectField(jObj, object_rect, obj_rect);

        env->CallBooleanMethod(jObjects, arraylist_add_method, jObj);
        env->DeleteLocalRef(jObj);
        env->DeleteLocalRef(obj_rect);
    }
    return 0;
}

extern "C" JNIEXPORT void

JNICALL
Java_com_lanyeo_liver_util_NCNN_resourceTaskDestroy(
        JNIEnv *env,
        jobject /* this */,
        jlong jTaskObjAddress
) {
    if (jTaskObjAddress == 0) {
        return;
    }
    auto* task = (NCNNInstance*)jTaskObjAddress;
    task->Destroy();
}